{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reaction MolGraph featurizers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chemprop.featurizers.molgraph.reaction import CondensedGraphOfReactionFeaturizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is an example reaction to featurize. The sanitizing code is to preserve atom mapped hydrogens in the graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit import Chem\n",
    "\n",
    "rct = Chem.MolFromSmiles(\"[H:1][C:4]([H:2])([H:3])[F:5]\", sanitize=False)\n",
    "pdt = Chem.MolFromSmiles(\"[H:1][C:4]([H:2])([H:3]).[F:5]\", sanitize=False)\n",
    "Chem.SanitizeMol(\n",
    "    rct, sanitizeOps=Chem.SanitizeFlags.SANITIZE_ALL ^ Chem.SanitizeFlags.SANITIZE_ADJUSTHS\n",
    ")\n",
    "Chem.SanitizeMol(\n",
    "    pdt, sanitizeOps=Chem.SanitizeFlags.SANITIZE_ALL ^ Chem.SanitizeFlags.SANITIZE_ADJUSTHS\n",
    ")\n",
    "\n",
    "rxn = (rct, pdt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Condensed Graph of Reaction featurizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like a [molecule](./molgraph_molecule_featurizer.ipynb) MolGraph featurizer, reaction MolGraph featurizers produce a `MolGraph`. The difference between the molecule and reaction versions is that a reaction takes two `rdkit.Chem.Mol` objects and need to know what \"mode\" of featurization to use. Available modes are found in `RxnMode`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reac_prod\n",
      "reac_prod_balance\n",
      "reac_diff\n",
      "reac_diff_balance\n",
      "prod_diff\n",
      "prod_diff_balance\n"
     ]
    }
   ],
   "source": [
    "from chemprop.featurizers import RxnMode\n",
    "\n",
    "for mode in RxnMode:\n",
    "    print(mode)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Briefly, \"reac\" stands for reactant features, \"prod\" stands for product features, and \"diff\" stands for the difference between reactant and product features. The two sets of features are concatenated together. \"balance\" refers to balancing imablanced reactions. See the 2022 [paper](https://doi.org/10.1021/acs.jcim.1c00975) by Heid and Green for more details. \"reac_diff\" is the default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "reac_diff = CondensedGraphOfReactionFeaturizer()\n",
    "reac_prod = CondensedGraphOfReactionFeaturizer(mode_=\"reac_prod\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  1, -1,\n",
       "         0,  0,  0,  0,  0, -1,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  1, -1,\n",
       "         0,  0,  0,  0,  0, -1,  0,  0,  0,  0,  0,  0]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reac_diff(rxn).E"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1,\n",
       "        0, 0, 0, 0, 0, 0],\n",
       "       [0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1,\n",
       "        0, 0, 0, 0, 0, 0],\n",
       "       [0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1,\n",
       "        0, 0, 0, 0, 0, 0],\n",
       "       [0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1,\n",
       "        0, 0, 0, 0, 0, 0],\n",
       "       [0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1,\n",
       "        0, 0, 0, 0, 0, 0],\n",
       "       [0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1,\n",
       "        0, 0, 0, 0, 0, 0],\n",
       "       [0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0],\n",
       "       [0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reac_prod(rxn).E"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like molecule MolGraph featurizers, reaction featurizers can use custom atom and bond featurizers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MolGraph(V=array([[ 1.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         1.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.01008,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ],\n",
       "       [ 0.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  1.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         1.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  1.     ,  0.     ,\n",
       "         0.     ,  0.12011,  0.     ,  0.     ,  0.     ,  1.     ,\n",
       "        -1.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ],\n",
       "       [ 1.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         1.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.01008,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ],\n",
       "       [ 1.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         1.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.01008,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ],\n",
       "       [ 0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  1.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         1.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  1.     ,  0.     ,\n",
       "         0.     ,  0.18998,  1.     , -1.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ]]), E=array([[ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  1, -1,\n",
       "         0,  0,  0,  0,  0, -1,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  1, -1,\n",
       "         0,  0,  0,  0,  0, -1,  0,  0,  0,  0,  0,  0]]), edge_index=array([[0, 1, 1, 2, 1, 3, 1, 4],\n",
       "       [1, 0, 2, 1, 3, 1, 4, 1]]), rev_edge_index=array([1, 0, 3, 2, 5, 4, 7, 6]))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from chemprop.featurizers import MultiHotAtomFeaturizer\n",
    "\n",
    "atom_featurizer = MultiHotAtomFeaturizer.organic()\n",
    "rxn_featurizer = CondensedGraphOfReactionFeaturizer(atom_featurizer=atom_featurizer)\n",
    "rxn_featurizer(rxn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Extra atom and bond features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Extra atom and bond features are not yet supported for reactions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "'atom_features_extra' is currently unsupported for reactions\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "MolGraph(V=array([[ 1.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         1.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.01008,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ],\n",
       "       [ 0.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  1.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         1.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  1.     ,  0.     ,\n",
       "         0.     ,  0.12011,  0.     ,  0.     ,  0.     ,  1.     ,\n",
       "        -1.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ],\n",
       "       [ 1.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         1.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.01008,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ],\n",
       "       [ 1.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         1.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.01008,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ],\n",
       "       [ 0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  1.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         1.     ,  0.     ,  1.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  1.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  1.     ,  0.     ,\n",
       "         0.     ,  0.18998,  1.     , -1.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ,  0.     ,  0.     ,  0.     ,\n",
       "         0.     ,  0.     ,  0.     ]]), E=array([[ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  1, -1,\n",
       "         0,  0,  0,  0,  0, -1,  0,  0,  0,  0,  0,  0],\n",
       "       [ 0,  1,  0,  0,  0,  0,  0,  1,  0,  0,  0,  0,  0,  0,  1, -1,\n",
       "         0,  0,  0,  0,  0, -1,  0,  0,  0,  0,  0,  0]]), edge_index=array([[0, 1, 1, 2, 1, 3, 1, 4],\n",
       "       [1, 0, 2, 1, 3, 1, 4, 1]]), rev_edge_index=array([1, 0, 3, 2, 5, 4, 7, 6]))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rxn_featurizer(rxn, atom_features_extra=[1.0])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
